{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Adding Models to NeuralForecast\n",
    "> Tutorial on how to add new models to NeuralForecast"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{.callout-warning}\n",
    "\n",
    "## Prerequisites\n",
    "\n",
    "This Guide assumes advanced familiarity with NeuralForecast.\n",
    "\n",
    "We highly recommend reading first the Getting Started and the NeuralForecast Map tutorials!\n",
    "\n",
    "Additionally, refer to the [CONTRIBUTING guide](https://github.com/Nixtla/neuralforecast/blob/main/CONTRIBUTING.md) for the basics how to contribute to NeuralForecast.\n",
    "\n",
    ":::"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This tutorial is aimed at contributors who want to add a new model to the NeuralForecast library. The library's existing modules handle optimization, training, selection, and evaluation of deep learning models. The `core` class simplifies building entire pipelines, both for industry and academia, on any dataset, with user-friendly methods such as `fit` and `predict`.\n",
    "\n",
    "<h4 style=\"text-align: center;\"> Adding a new model to NeuralForecast is simpler than building a new PyTorch model from scratch. You only need to write the forward method. </h4>\n",
    "\n",
    "**It has the following additional advantages:**\n",
    "\n",
    "* Existing modules in NeuralForecast already implement the essential training and evaluating aspects for deep learning models.\n",
    "* Integrated with PyTorch-Lightning and Tune libraries for efficient optimization and distributed computation.\n",
    "* The `BaseModel` classes provide common optimization components, such as early stopping and learning rate schedulers.\n",
    "* Automatic performance tests are scheduled on Github to ensure quality standards.\n",
    "* Users can easily compare the performance and computation of the new model with existing models.\n",
    "* Opportunity for exposure to a large community of users and contributors."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example: simplified MLP model\n",
    "\n",
    "We will present the tutorial following an example on how to add a simplified version of the current `MLP` model, which does not include exogenous covariates.\n",
    "\n",
    "At a given timestamp $t$, the `MLP` model will forecast the next $h$ values of the univariate target time, $Y_{t+1:t+h}$, using as inputs the last $L$ historical values, given by $Y_{t-L:t}$. The following figure presents a diagram of the model."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Figure 1. Three layer MLP with autoregresive inputs.](../../imgs_models/mlp.png)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Preliminaries\n",
    "\n",
    "Follow our tutorial on contributing [here](https://github.com/Nixtla/neuralforecast/blob/main/CONTRIBUTING.md) to set up your development environment.\n",
    "\n",
    "Here is a short list of the most important steps:\n",
    "\n",
    "1. Create a fork of the `neuralforecast` library.\n",
    "2. Clone the fork to your computer.\n",
    "3. Set an environment with the `neuralforecast` library, core dependencies, and `nbdev` package to code your model in an interactive notebook."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Inherit the Base Class (`BaseModel`)\n",
    "\n",
    "The library contains a base model class: `BaseModel`. Using class attributes we can make this model recurrent or not, or multivariate or univariate, or allow the use of exogenous inputs.\n",
    "\n",
    "### a. Sampling process\n",
    "\n",
    "During training, the base class receives a sample of time series of the dataset from the `TimeSeriesLoader` module. The `BaseModel` models will sample individual windows of size `input_size+h`, starting from random timestamps.\n",
    "\n",
    "### b. `BaseModel`' hyperparameters\n",
    "\n",
    "Get familiar with the hyperparameters specified in the base class, including `h` (horizon), `input_size`, and optimization hyperparameters such as `learning_rate`, `max_steps`, among others. The following list presents the hyperparameters related to the sampling of windows:\n",
    " \n",
    " * `h` (h): number of future values to predict.\n",
    " * `input_size` (L): number of historic values to use as input for the model.\n",
    " * `batch_size` (bs): number of time series sampled by the loader during training.\n",
    " * `valid_batch_size` (v_bs): number of time series sampled by the loader during inference (validation and test).\n",
    " * `windows_batch_size` (w_bs): number of individual windows sampled during training (from the previous time series) to form the batch.\n",
    " * `inference_windows_batch_size` (i_bs): number of individual windows sampled during inference to form each batch. Used to control the GPU memory.\n",
    "\n",
    "### c. Input and Output batch shapes\n",
    "\n",
    "The `forward` method receives a batch of data in a dictionary with the following keys:\n",
    "\n",
    "- `insample_y`: historic values of the time series.\n",
    "- `insample_mask`: mask indicating the available values of the time series (1 if available, 0 if missing).\n",
    "- `futr_exog`: future exogenous covariates (if any).\n",
    "- `hist_exog`: historic exogenous covariates (if any).\n",
    "- `stat_exog`: static exogenous covariates (if any).\n",
    "\n",
    "The following table presents the shape for each tensor if the attribute `MULTIVARIATE = False` is set:\n",
    "\n",
    "| `tensor`        | `BaseModel`            |\n",
    "|-----------------|--------------------------|\n",
    "| `insample_y`    | (`w_bs`, `L`, `1`)       |\n",
    "| `insample_mask` | (`w_bs`, `L`)            |\n",
    "| `futr_exog`     | (`w_bs`, `L`+`h`, `n_f`) |\n",
    "| `hist_exog`     | (`w_bs`, `L`, `n_h`)     |\n",
    "| `stat_exog`     | (`w_bs`,`n_s`)           |\n",
    "\n",
    "The `forward` function should return a single tensor with the forecasts of the next `h` timestamps for each window. Use the attributes of the `loss` class to automatically parse the output to the correct shape (see the example below).  \n",
    "\n",
    ":::{.callout-tip}\n",
    "\n",
    "Since we are using `nbdev`, you can easily add prints to the code and see the shapes of the tensors during training.\n",
    "\n",
    ":::\n",
    "\n",
    "### d. `BaseModel`' methods\n",
    "\n",
    "The `BaseModel` class contains several common methods for all windows-based models, simplifying the development of new models by preventing code duplication. The most important methods of the class are:\n",
    "\n",
    "* `_create_windows`: parses the time series from the `TimeSeriesLoader` into individual windows of size `input_size+h`.\n",
    "* `_normalization`: normalizes each window based on the `scaler` type.\n",
    "* `_inv_normalization`: inverse normalization of the forecasts.\n",
    "* `training_step`: training step of the model, called by PyTorch-Lightning's `Trainer` class during training (`fit` method).\n",
    "* `validation_step`: validation step of the model, called by PyTorch-Lightning's `Trainer` class during validation.\n",
    "* `predict_step`: prediction step of the model, called by PyTorch-Lightning's `Trainer` class during inference (`predict` method)."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Create the model file and class\n",
    "\n",
    "Once familiar with the basics of the `BaseModel` class, the next step is creating your particular model.\n",
    "\n",
    "The main steps are:\n",
    "\n",
    "1. Create the file in the `nbs` folder (https://github.com/Nixtla/neuralforecast/tree/main/nbs). It should be named `models.YOUR_MODEL_NAME.ipynb`.\n",
    "2. Add the header of the `nbdev` file.\n",
    "3. Import libraries in the file. \n",
    "4. Define the `__init__` method with the model's inherited and particular hyperparameters and instantiate the architecture.\n",
    "5. Set the following model attributes:\n",
    "    - `EXOGENOUS_FUTR`: if the model can handle future exogenous variables (True) or not (False)\n",
    "    - `EXOGENOUS_HIST`: if the model can handle historical exogenous variables (True) or not (False)\n",
    "    - `EXOGENOUS_STAT`: if the model can handle static exogenous variables (True) or not (False)\n",
    "    - `MULTIVARIATE`: If the model produces multivariate forecasts (True) or univariate (False)\n",
    "    - `RECURRENT`: If the model produces forecasts recursively (True) or direct (False)\n",
    "5. Define the `forward` method, which recieves the input batch dictionary and returns the forecast."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### a. Model class"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, add the following **two cells** on top of the `nbdev` file.\n",
    "\n",
    "```python\n",
    "#| default_exp models.mlp\n",
    "```\n",
    "\n",
    ":::{.callout-important}\n",
    "\n",
    "Change `mlp` to your model's name, using lowercase and underscores. When you later run `nbdev_export`, it will create a `YOUR_MODEL.py` script in the `neuralforecast/models/` directory.\n",
    "\n",
    ":::\n",
    "\n",
    "```python\n",
    "#| echo: false\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "```\n",
    "\n",
    "Next, add the dependencies of the model.\n",
    "\n",
    "```python\n",
    "#| export\n",
    "from typing import Optional\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "from neuralforecast.losses.pytorch import MAE\n",
    "from neuralforecast.common._base_model import BaseModel\n",
    "```\n",
    "\n",
    ":::{.callout-tip}\n",
    "\n",
    "Don't forget to add the `#| export` tag on this cell.\n",
    "\n",
    ":::"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, create the class with the `init` and `forward` methods. The following example shows the example for the simplified `MLP` model. We explain important details after the code."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```python\n",
    "#| export\n",
    "class MLP(BaseModel): # <<---- Inherits from BaseModel\n",
    "    # Set class attributes to determine this model's characteristics\n",
    "    EXOGENOUS_FUTR = False   # If the model can handle future exogenous variables\n",
    "    EXOGENOUS_HIST = False   # If the model can handle historical exogenous variables\n",
    "    EXOGENOUS_STAT = False   # If the model can handle static exogenous variables\n",
    "    MULTIVARIATE = False    # If the model produces multivariate forecasts (True) or univariate (False)\n",
    "    RECURRENT = False       # If the model produces forecasts recursively (True) or direct (False)\n",
    "\n",
    "    def __init__(self,\n",
    "                 # Inhereted hyperparameters with no defaults\n",
    "                 h,\n",
    "                 input_size,\n",
    "                 # Model specific hyperparameters\n",
    "                 num_layers = 2,\n",
    "                 hidden_size = 1024,\n",
    "                 # Inhereted hyperparameters with defaults\n",
    "                 futr_exog_list = None,\n",
    "                 hist_exog_list = None,\n",
    "                 stat_exog_list = None,                 \n",
    "                 exclude_insample_y = False,\n",
    "                 loss = MAE(),\n",
    "                 valid_loss = None,\n",
    "                 max_steps: int = 1000,\n",
    "                 learning_rate: float = 1e-3,\n",
    "                 num_lr_decays: int = -1,\n",
    "                 early_stop_patience_steps: int =-1,\n",
    "                 val_check_steps: int = 100,\n",
    "                 batch_size: int = 32,\n",
    "                 valid_batch_size: Optional[int] = None,\n",
    "                 windows_batch_size = 1024,\n",
    "                 inference_windows_batch_size = -1,\n",
    "                 start_padding_enabled = False,\n",
    "                 step_size: int = 1,\n",
    "                 scaler_type: str = 'identity',\n",
    "                 random_seed: int = 1,\n",
    "                 drop_last_loader: bool = False,\n",
    "                 optimizer = None,\n",
    "                 optimizer_kwargs = None,\n",
    "                 lr_scheduler = None,\n",
    "                 lr_scheduler_kwargs = None,\n",
    "                 dataloader_kwargs = None,\n",
    "                 **trainer_kwargs):\n",
    "    # Inherit BaseWindows class\n",
    "    super(MLP, self).__init__(h=h,\n",
    "                              input_size=input_size,\n",
    "                              ..., # <<--- Add all inhereted hyperparameters\n",
    "                              random_seed=random_seed,\n",
    "                              **trainer_kwargs)\n",
    "\n",
    "    # Architecture\n",
    "    self.num_layers = num_layers\n",
    "    self.hidden_size = hidden_size\n",
    "\n",
    "    # MultiLayer Perceptron\n",
    "    layers = [nn.Linear(in_features=input_size, out_features=hidden_size)]\n",
    "    layers += [nn.ReLU()]\n",
    "    for i in range(num_layers - 1):\n",
    "        layers += [nn.Linear(in_features=hidden_size, out_features=hidden_size)]\n",
    "        layers += [nn.ReLU()]\n",
    "    self.mlp = nn.ModuleList(layers)\n",
    "\n",
    "    # Adapter with Loss dependent dimensions\n",
    "    self.out = nn.Linear(in_features=hidden_size, \n",
    "                         out_features=h * self.loss.outputsize_multiplier) ## <<--- Use outputsize_multiplier to adjust output size\n",
    "\n",
    "    def forward(self, windows_batch): # <<--- Receives windows_batch dictionary\n",
    "        # Parse windows_batch\n",
    "        insample_y = windows_batch['insample_y'].squeeze(-1)                            # [batch_size, input_size]\n",
    "        # MLP\n",
    "        hidden = self.mlp(insample_y)                                                   # [batch_size, hidden_size]\n",
    "        y_pred = self.out(hidden)                                                       # [batch_size, h * n_outputs]\n",
    "        \n",
    "        # Reshape\n",
    "        y_pred = y_pred.reshape(batch_size, self.h, self.loss.outputsize_multiplier)    # [batch_size, h, n_outputs]\n",
    "\n",
    "        return y_pred\n",
    "\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ":::{.callout-tip}\n",
    "\n",
    "* Don't forget to add the `#| export` tag on each cell.\n",
    "* Larger architectures, such as Transformers, might require splitting the `forward` by using intermediate functions.\n",
    "\n",
    ":::"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Important notes\n",
    "\n",
    "The base class has many hyperparameters, and models must have default values for all of them (except `h` and `input_size`). If you are unsure of what default value to use, we recommend copying the default values from existing models for most optimization and sampling hyperparameters. You can change the default values later at any time.\n",
    "\n",
    "The `reshape` method at the end of the `forward` step is used to adjust the output shape. The `loss` class contains an `outputsize_multiplier` attribute to automatically adjust the output size of the forecast depending on the `loss`. For example, for the Multi-quantile loss (`MQLoss`), the model needs to output each quantile for each horizon.\n",
    "\n",
    "### b. Tests and documentation\n",
    "\n",
    "`nbdev` allows for testing and documenting the model during the development process. It allows users to iterate the development within the notebook, testing the code in the same environment. Refer to existing models, such as the complete MLP model [here](https://github.com/Nixtla/neuralforecast/blob/main/nbs/models.mlp.ipynb). These files already contain the tests, documentation, and usage examples that were used during the development process.\n",
    "\n",
    "### c. Export the new model to the library with `nbdev`\n",
    "\n",
    "Following the CONTRIBUTING guide, the next step is to export the new model from the development notebook to the `neuralforecast` folder with the actual scripts.\n",
    "\n",
    "To export the model, run `nbdev_export` in your terminal. You should see a new file with your model in the `neuralforecast/models/` folder."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Core class and additional files\n",
    "\n",
    "Finally, add the model to the `core` class and additional files:\n",
    "\n",
    "1. Manually add the model in the following [init file](https://github.com/Nixtla/neuralforecast/blob/main/neuralforecast/models/__init__.py).\n",
    "2. Add the model to the `core` class, using the `nbdev` file [here](https://github.com/Nixtla/neuralforecast/blob/main/nbs/core.ipynb):\n",
    "    \n",
    "    a. Add the model to the initial model list:\n",
    "    ```python\n",
    "    from neuralforecast.models import (\n",
    "    GRU, LSTM, RNN, TCN, DilatedRNN,\n",
    "    MLP, NHITS, NBEATS, NBEATSx,\n",
    "    TFT, VanillaTransformer,\n",
    "    Informer, Autoformer, FEDformer,\n",
    "    StemGNN, PatchTST\n",
    "    )\n",
    "    ```\n",
    "    b. Add the model to the `MODEL_FILENAME_DICT` dictionary (used for the `save` and `load` functions)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Add the model to the documentation\n",
    "\n",
    "It's important to add the model to the necessary documentation pages so that everyone can find the documentation:\n",
    "\n",
    "1. Add the model to the [model overview table](https://github.com/Nixtla/neuralforecast/blob/main/nbs/docs/capabilities/01_overview.ipynb).\n",
    "2. Add the model to the [sidebar](https://github.com/Nixtla/neuralforecast/blob/main/nbs/sidebar.yml) for the API reference.\n",
    "3. Add the model to [mint.json](https://github.com/Nixtla/neuralforecast/blob/main/nbs/mint.json)."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Upload to GitHub\n",
    "\n",
    "Congratulations! The model is ready to be used in the library following the steps above. \n",
    "\n",
    "Follow our contributing guide's final steps to upload the model to GitHub: [here](https://github.com/Nixtla/neuralforecast/blob/main/CONTRIBUTING.md).\n",
    "\n",
    "One of the maintainers will review the PR, request changes if necessary, and merge it into the library."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Quick Checklist\n",
    "\n",
    "* Get familiar with the `BaseModel` class hyperparameters and input/output shapes of the `forward` method.\n",
    "* Create the notebook with your model class in the `nbs` folder: `models.YOUR_MODEL_NAME.ipynb`\n",
    "* Add the header and import libraries.\n",
    "* Implement `init` and `forward` methods and set the class attributes.\n",
    "* Export model with `nbdev_export`.\n",
    "* Add model to this [init file](https://github.com/Nixtla/neuralforecast/blob/main/neuralforecast/models/__init__.py).\n",
    "* Add the model to the `core ` class [here](https://github.com/Nixtla/neuralforecast/blob/main/nbs/core.ipynb).\n",
    "* Follow the CONTRIBUTING guide to create the PR to upload the model.\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
